import sys
import socket
import time
import curses
import matplotlib.pyplot as plt
import matplotlib as mpl
import numpy as np

vmin = 1000
vmax = 1500


############################################################
def cxpd_cmd_gen(index,arg5,arg4,arg3,arg2,arg1,arg0):
	cmd_buf = bytearray(8)
	cmd_buf[0] = 0x00
	cmd_buf[1] = index
	cmd_buf[2] = arg5
	cmd_buf[3] = arg4
	cmd_buf[4] = arg3
	cmd_buf[5] = arg2
	cmd_buf[6] = arg1
	cmd_buf[7] = arg0
	return cmd_buf

def list_to_buf(list_in):
	buf_out = bytearray(0)
	for i in list_in :
		for j in i :
			buf_out.append(j)
	return buf_out

############################################################

def data_tx(tcp_client_socket,cmd_start_buf,frame_num):
    for _ in range(frame_num):
        time.sleep(0.02)
        tcp_client_socket.sendall(cmd_start_buf)


def data_rx_muti(tcp_client_socket,buf_size):
    cmd_tcp_rst = bytearray(8)
    cmd_tcp_rst[1] = 0x04
    buf = bytearray(buf_size)
    view_buf = memoryview(buf)
    toread = buf_size
    view_tmp = view_buf
    while toread:
        nbytes = tcp_client_socket.recv_into(view_tmp, toread)
        view_tmp = view_tmp[nbytes:]
        toread -= nbytes
    return buf


def main(stdscr):
    stdscr.nodelay(True)
    print("PC Process Start\r")
    if len(sys.argv) > 1:
        frame_num = int(sys.argv[1])
        if frame_num > 0:
            pass
        else:
            frame_num = -1
    else:
        frame_num = -1

    if len(sys.argv) > 2:
        a_rst_in = int(sys.argv[2])
    else:
        a_rst_in = 780
    a_rst = int(a_rst_in/2500*65536)
    a_rst_d0 = int((a_rst%16)*16)
    a_rst_d1 = int((a_rst%4096)>>4)
    a_rst_d2 = int(a_rst>>12)+0x10

    if len(sys.argv) > 3:
        delay_set = int(sys.argv[3])
    else :
        delay_set = 3

    # /////////////////////////////////////////////////////////////////////////////
    # VOLTAGE initlal
    cmd_initial_list = []
    cmd_initial_list.append(cxpd_cmd_gen(0x01,0x00,0x00,0x0a,0x00,0x00,0x00))
    cmd_initial_list.append(cxpd_cmd_gen(0x01,0x00,0x00,0x07,0x00,0x00,0x00))
    cmd_initial_list.append(cxpd_cmd_gen(0x01,0x00,0x00,0x09,0x0a,0x00,0x00))
    cmd_initial_list.append(cxpd_cmd_gen(0x01,0x00,0x00,0x04,0x00,0x00,0xff))
    # TOPMETAL II voltage set
    cmd_initial_list.append(cxpd_cmd_gen(0x01,0x00,0x00,0x03,0x03,0xff,0xf0))# 630mV on csa_ref
    if a_rst_in < 700 or a_rst_in > 1200:
        cmd_initial_list.append(cxpd_cmd_gen(0x01,0x00,0x00,0x03,0x14,0xfd,0xf0))# 780mV on a_rst
    else:
        cmd_initial_list.append(cxpd_cmd_gen(0x01,0x00,0x00,0x03,a_rst_d2,a_rst_d1,a_rst_d0))# user set a_rst
    cmd_initial_list.append(cxpd_cmd_gen(0x01,0x00,0x00,0x03,0x27,0x46,0xd0))# 1137mV on ad_offset
    cmd_initial_list.append(cxpd_cmd_gen(0x03,0xff,0xff,0xff,0xff,0xff,0xff))
    cmd_initial_list.append(cxpd_cmd_gen(0x03,0xff,0xff,0xff,0xff,0xff,0xff))

    cmd_start_list = []
    cmd_start_list.append(cxpd_cmd_gen(0x02,0x00,0x00,0x00,0x00,0x00,0x01))

    cmd_initial_buf = list_to_buf(cmd_initial_list)
    cmd_start_buf = list_to_buf(cmd_start_list)

    # connect to FPGA
    tcp_client_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    try:
        print("connecting to FPGA ... \r")
        tcp_client_socket.connect(("192.168.1.103",1024))
        print('TCP connected\r')
    except socket.error:
        print("fail to connect  to FPGA\r")
        exit(0)

    
    tcp_client_socket.sendall(cmd_initial_buf)

    tcp_buf_size = 5184 * 2

    for i in range(delay_set):
        time.sleep(1)
        print('ready to read data : ',delay_set-i, '\r')

    fig, ax = plt.subplots()
    norm = mpl.colors.Normalize(vmin=vmin,vmax=vmax)
    i = 0
    while frame_num == -1 or i < frame_num:
        data_tx(tcp_client_socket, cmd_start_buf, 1)
        data_buf = data_rx_muti(tcp_client_socket,tcp_buf_size)
        data = np.frombuffer(data_buf,dtype='>u2')
        data = data.reshape(5184)
        ax.cla()
        data_deal = data.reshape(72,72)
        ax.imshow(data_deal)
        ax.set_title("frame {}".format(i))
        pcm = ax.pcolormesh(data_deal,cmap='viridis',norm=norm)
        if i == 0:
            fig.colorbar(pcm, ax=ax)
        plt.pause(0.01)
        ch = stdscr.getch()
        if ch != -1:
             break
        i += 1

    tcp_client_socket.close()
    plt.close()

    print('Monitoring finished. Goodbye.')


############################################################
# PROCESS START #
if __name__ == '__main__':
    curses.wrapper(main)
